#include <string>
#include <thread>
#include <stdio.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include "hilog/log.h"
#include <netdb.h>
#include <netinet/in.h>
#include <huks/native_huks_api.h>
#include <huks/native_huks_param.h>
#include <huks/native_huks_type.h>
#include <map>
#include <vector>
#include <algorithm>
#include <stdlib.h>
#include <string.h> 
#include "leaf.h"
#include "json11/json11.hpp"

using namespace json11;

#define logger_info(fmt, args...)  OH_LOG_Print(LOG_APP, LOG_ERROR, 0x1234, "debug", fmt, ##args)
#define BUFFER_SIZE 2048

const std::string ConfigTemplate = R"delimiter(
{
    "log": {
        "level": "info"
    },
    "dns": {
        "servers": [
            "114.114.114.114",
            "223.5.5.5"
        ]
    },
    "inbounds": [
        {
            "protocol": "tun",
            "settings": {
                "fd": 0
            },
            "tag": "tun_in"
        }
    ],
    "outbounds": [
        {
            "bind": "10.0.0.1",
            "protocol": "socks",
            "settings": {
                "address": "127.0.0.1",
                "port": 30038
            },
            "tag": "socks_out"
        }
    ],
    "rules": [
    ]
}
)delimiter";

std::string get_config_json(int fd, const char *host, int port) {
    std::string err;
    const auto json = Json::parse(ConfigTemplate, err);
    json11::Json::object jsonObjectMap = json.object_items();
    std::vector<json11::Json> inbounds = jsonObjectMap["inbounds"].array_items();
    json11::Json::object inbound0 = inbounds[0].object_items();
    json11::Json::object settings = inbound0["settings"].object_items();
    settings["fd"] = fd;
    inbound0["settings"] = settings;
    inbounds[0] = inbound0;
    jsonObjectMap["inbounds"] = inbounds;

    std::vector<json11::Json> outbounds = jsonObjectMap["outbounds"].array_items();
    json11::Json::object outbound0 = outbounds[0].object_items();
    settings = outbound0["settings"].object_items();
    settings["address"] = std::string(host);
    settings["port"] = port;
    outbound0["settings"] = settings;
    outbounds[0] = outbound0;
    jsonObjectMap["outbounds"] = outbounds;

    json11::Json newJsonObject = json11::Json(jsonObjectMap);

    auto json_str = newJsonObject.dump();
    return json_str;
}

extern "C" int leaf_run_with_config_string(uint16_t rt_id, const char *config);

void start_leaf(int fd, const char *proxy_host, int proxy_port) {
    auto json_str = get_config_json(fd, proxy_host, proxy_port);
    leaf_run_with_config_string(0, json_str.c_str());
}

static int stack_fd = 0;
static std::mutex g_tun_mutex;
// DNS查询上下文
struct DnsQueryContext {
    std::string domain;
    uint32_t ip;
    dns_query_status_t status;
    std::mutex mutex;
    std::condition_variable cv;
};
// 本地DNS缓存
static std::map<std::string, uint32_t> dns_local_cache;
static std::mutex dns_cache_mutex;
static std::string dns_upstream = "8.8.8.8";
static dns_resolve_cb_t dns_callback = nullptr;
static void* dns_callback_userdata = nullptr;
// 多网关池
static std::vector<Gateway> gateway_pool;
static gateway_change_cb_t gateway_callback = nullptr;
static void* gateway_callback_userdata = nullptr;
// 资源URI映射表
static std::map<std::string, std::string> resource_map;
// HTTP会话管理
static http_session_cb_t http_session_callback = nullptr;
static void* http_callback_userdata = nullptr;
static bool http_analysis_enabled = false;
static std::string http_secret_key = "default-secret-key";
// 安全密钥存储（实际项目应使用HUKS安全存储）
static const char *g_httpHmacKey = "your_secure_hmac_key_256bit";
// HMAC参数配置
static struct OH_Huks_Param g_hmacParams[] = {
    {
        .tag = OH_HUKS_TAG_ALGORITHM,
        .uint32Param = OH_HUKS_ALG_HMAC
    }, {
        .tag = OH_HUKS_TAG_PURPOSE,
        .uint32Param = OH_HUKS_KEY_PURPOSE_MAC
    }, {
        .tag = OH_HUKS_TAG_KEY_SIZE,
        .uint32Param = OH_HUKS_AES_KEY_SIZE_256 // HMAC-SHA256推荐256位密钥
    }, {
        .tag = OH_HUKS_TAG_DIGEST,
        .uint32Param = OH_HUKS_DIGEST_SHA256
    }
};
// PacketProtocol回调变量
static packet_callback_t g_packet_callback = nullptr;
static void* g_packet_callback_userdata = nullptr;

void set_packet_callback(packet_callback_t cb, void* userdata) {
    g_packet_callback = cb;
    g_packet_callback_userdata = userdata;
}

// 初始化参数设置
OH_Huks_Result InitParamSet(struct OH_Huks_ParamSet **paramSet, const struct OH_Huks_Param *params, uint32_t paramCount)
{
    OH_Huks_Result ret = OH_Huks_InitParamSet(paramSet);
    if (ret.errorCode != OH_HUKS_SUCCESS) {
        return ret;
    }
    ret = OH_Huks_AddParams(*paramSet, params, paramCount);
    if (ret.errorCode != OH_HUKS_SUCCESS) {
        OH_Huks_FreeParamSet(paramSet);
        return ret;
    }
    ret = OH_Huks_BuildParamSet(paramSet);
    if (ret.errorCode != OH_HUKS_SUCCESS) {
        OH_Huks_FreeParamSet(paramSet);
        return ret;
    }
    return ret;
}

// 基于会话模式的HMAC计算
OH_Huks_Result ComputeHmac(const struct OH_Huks_Blob *keyAlias, const struct OH_Huks_ParamSet *paramSet, const struct OH_Huks_Blob *inData, struct OH_Huks_Blob *outData) 
{
    uint8_t handleBuf[sizeof(uint64_t)] = {0};
    struct OH_Huks_Blob handle = {sizeof(uint64_t), handleBuf};

    // 1. 初始化会话
    OH_Huks_Result ret = OH_Huks_InitSession(keyAlias, paramSet, &handle, nullptr);
    if (ret.errorCode != OH_HUKS_SUCCESS) {
        logger_info("zjx--InitSession failed! errorCode=%{public}d", ret.errorCode);
        return ret;
    }

    // 2. 单次完成计算（适合HTTP签名场景）
    ret = OH_Huks_FinishSession(&handle, paramSet, inData, outData);
    if (ret.errorCode != OH_HUKS_SUCCESS) {
        logger_info("zjx--FinishSession failed! errorCode=%{public}d", ret.errorCode);
    }

    return ret;
}

//http会话回调函数
void http_set_session_callback(http_session_cb_t cb, void* userdata) {
    http_session_callback = cb;
    http_callback_userdata = userdata;
}
// http协议开关
void http_enable_protocol_analysis(bool enable) {
    http_analysis_enabled = enable;
}

// HTTP签名方法
uint32_t http_generate_signature(const char* data, size_t len) {
    /* 1. 准备密钥材料 */
    struct OH_Huks_Blob keyAlias = {
        (uint32_t)strlen("http_hmac_key"),
        (uint8_t *)"http_hmac_key"
    };
    struct OH_Huks_Blob keyMaterial = {
        (uint32_t)strlen(g_httpHmacKey),
        (uint8_t *)g_httpHmacKey
    };
    /* 2. 准备输入数据 */
    struct OH_Huks_Blob inData = {
        (uint32_t)len,
        (uint8_t *)data
    };
    /* 3. 初始化参数集 */
    struct OH_Huks_ParamSet *paramSet = nullptr;
    OH_Huks_Result ohResult;
    ohResult = InitParamSet(&paramSet, g_hmacParams, sizeof(g_hmacParams)/sizeof(g_hmacParams[0]));
    if (ohResult.errorCode != OH_HUKS_SUCCESS) {
        logger_info("zjx--InitParamSet failed! errorCode=%{public}d", ohResult.errorCode);
        return 0;
    }
    /* 4. 导入密钥（生产环境应预生成） */
    ohResult = OH_Huks_ImportKeyItem(&keyAlias, paramSet, &keyMaterial);
    if (ohResult.errorCode != OH_HUKS_SUCCESS) {
        logger_info("zjx--ImportKey failed! errorCode=%{public}d", ohResult.errorCode);
        OH_Huks_FreeParamSet(&paramSet);
        return 0;
    }
    /* 5. 计算HMAC */
    uint8_t hmacResult[OH_HUKS_DIGEST_SHA256] = {0};
    struct OH_Huks_Blob outData = {sizeof(hmacResult), hmacResult};

    ohResult = ComputeHmac(&keyAlias, paramSet, &inData, &outData);
    /* 6. 清理资源 */
    OH_Huks_DeleteKeyItem(&keyAlias, paramSet);
    OH_Huks_FreeParamSet(&paramSet);

    if (ohResult.errorCode != OH_HUKS_SUCCESS) {
        return 0;
    }
    /* 7. 生成4字节签名 */
    uint32_t signature = 0;
    if (memcpy(&signature, hmacResult, sizeof(signature)) != 0) {
        logger_info("zjx--memcpy failed!");
        return 0;
    }

    return ntohl(signature);
}

// 解析HTTP请求
static bool parse_http_request(const uint8_t* data, size_t len, 
                             HttpSessionStamp* stamp) {
    if (len < 16) return false; // 最小HTTP请求长度
    
    const char* ptr = (const char*)data;
    const char* end = ptr + len;
    
    // 解析方法
    const char* space = strchr(ptr, ' ');
    if (!space || space - ptr >= 8) return false;
    memcpy(stamp->method, ptr, space - ptr);
    stamp->method[space - ptr] = '\0';
    
    // 解析URI
    ptr = space + 1;
    space = strchr(ptr, ' ');
    if (!space || space - ptr >= 512) return false;
    memcpy(stamp->uri, ptr, space - ptr);
    stamp->uri[space - ptr] = '\0';
    
    // 查找Host头
    ptr = strstr(ptr, "\r\nHost: ");
    if (!ptr) return false;
    ptr += 8;
    const char* host_end = strstr(ptr, "\r\n");
    if (!host_end || host_end - ptr >= 256) return false;
    memcpy(stamp->host, ptr, host_end - ptr);
    stamp->host[host_end - ptr] = '\0';
    
    return true;
}

int get_packet_protocol(const uint8_t* data, size_t len) {
    if (!data) {
        logger_info("zjx--Packet data is NULL");
        return -1;
    }
    
    if (len < 20) {
        logger_info("zjx--Packet too short: len=%{public}zu (need >=20)", len);
        return -1;
    }
    
    // 调试打印前20个字节
    logger_info("zjx--get Packet header (first 20 bytes):");
    for (size_t i = 0; i < (len > 20 ? 20 : len); ++i) {
        logger_info("zjx--get Packet header result :[%{public}02zu]: 0x%{public}02x", i, data[i]);
    }
    
    // 检查 IP (前 4 bits)
    uint8_t ver = data[0] >> 4;
    if (ver != 4) {
        logger_info("zjx--Not IPv4: version=%{public}d (expected 4)", ver);
        return -1; 
    }
    
    // 从字节9获取协议 (基于指数)
    uint8_t protocol = data[9];
    logger_info("zjx--Detected protocol: 0x%{public}02x", protocol);
    
    // 返回 TCP(6)/UDP(17), 其他 -1
    return (protocol == 6 || protocol == 17) ? protocol : -1;
}

// DNS查询线程函数
void dns_query_thread(DnsQueryContext* ctx) {
    struct addrinfo hints = {0}, *result = nullptr;
    hints.ai_family = AF_INET; // 只查询IPv4地址
    hints.ai_socktype = SOCK_DGRAM;
    hints.ai_flags = AI_PASSIVE;
    
    int ret = getaddrinfo(ctx->domain.c_str(), nullptr, &hints, &result);
    {
        std::lock_guard<std::mutex> lock(ctx->mutex);
        if (ret == 0 && result != nullptr) {
            struct sockaddr_in* addr = (struct sockaddr_in*)result->ai_addr;
            ctx->ip = addr->sin_addr.s_addr;
            ctx->status = DNS_QUERY_SUCCESS;
            
            // 添加到本地缓存
            {
                std::lock_guard<std::mutex> cache_lock(dns_cache_mutex);
                dns_local_cache[ctx->domain] = ctx->ip;
            }
            
            logger_info("zjx--DNS query success: %{public}s -> %{public}u", 
                       ctx->domain.c_str(), ctx->ip);
        } else {
            ctx->status = DNS_QUERY_FAILED;
            logger_info("zjx--DNS query failed: %{public}s, error: %{public}d", 
                       ctx->domain.c_str(), ret);
        }
    }
    ctx->cv.notify_one();
    
    if (result) {
        freeaddrinfo(result);
    }
    
    // 通知回调
    if (dns_callback) {
        dns_callback(ctx->domain.c_str(), ctx->ip, dns_callback_userdata);
    }
    
    delete ctx;
}

uint32_t dns_resolve_local(const char* domain) {
    // 首先检查本地缓存
    {
        std::lock_guard<std::mutex> lock(dns_cache_mutex);
        auto it = dns_local_cache.find(domain);
        if (it != dns_local_cache.end()) {
            logger_info("zjx--DNS cache hit: %{public}s -> %{public}u", domain, it->second);
            return it->second;
        }
    }
    
    logger_info("zjx--DNS cache miss, query upstream: %{public}s", domain);
    
    // 创建查询上下文
    DnsQueryContext* ctx = new DnsQueryContext();
    ctx->domain = domain;
    ctx->ip = 0;
    ctx->status = DNS_QUERY_IN_PROGRESS;
    
    // 启动查询线程
    std::thread(dns_query_thread, ctx).detach();
    
    // 等待查询完成或超时
    {
        std::unique_lock<std::mutex> lock(ctx->mutex);
        if (ctx->cv.wait_for(lock, std::chrono::seconds(3)) == std::cv_status::timeout) {
            logger_info("zjx--DNS query timeout: %{public}s", domain);
            ctx->status = DNS_QUERY_FAILED;
        }
    }
    
    uint32_t result_ip = ctx->ip;
    if (ctx->status != DNS_QUERY_SUCCESS) {
        result_ip = 0;
    }
    
    // 注意：ctx会在查询线程中被删除
    return result_ip;
}

//DNS功能实现
void dns_set_local(const char* domain, uint32_t ip) {
    std::lock_guard<std::mutex> lock(dns_cache_mutex);
    dns_local_cache[domain] = ip;
}

void dns_set_upstream(const char* dns_server) {
    dns_upstream = dns_server;
    
    // 设置系统DNS服务器（可选）
    // 注意：这可能需要root权限
    /*
    struct in_addr addr;
    if (inet_pton(AF_INET, dns_server, &addr) == 1) {
        res_init();
        _res.nscount = 1;
        _res.nsaddr_list[0].sin_addr = addr;
        _res.nsaddr_list[0].sin_family = AF_INET;
        _res.nsaddr_list[0].sin_port = htons(53);
    }
    */
}

void dns_set_callback(dns_resolve_cb_t cb, void* userdata) {
    dns_callback = cb;
    dns_callback_userdata = userdata;
}

void gateway_add(const Gateway* gw) {
    // 去重处理
    auto pos = std::find_if(gateway_pool.begin(), gateway_pool.end(),
        [&](const Gateway& x) { return strcmp(x.host, gw->host) == 0; });
    
    if (pos == gateway_pool.end()) {
        gateway_pool.push_back(*gw);
        // 通知新网关添加
        if (gateway_callback) {
            gateway_callback(gw->host, gw->port, gw->latency, gateway_callback_userdata);
        }
    }
}

void gateway_remove(const char* host) {
    if (!host) return;

    auto it = std::remove_if(gateway_pool.begin(), gateway_pool.end(),
        [&](const Gateway& x) { 
            return x.host && strcmp(x.host, host) == 0;
        });
    
    if (it != gateway_pool.end()) {
        gateway_pool.erase(it, gateway_pool.end());
        if (gateway_callback) {
            gateway_callback(host, 0, 0, gateway_callback_userdata);
        }
        logger_info("zjx--Gateway removed: %{public}s", host);
    }
}

void gateway_update_latency(const char* host, uint32_t latency) {
    for (auto& gw : gateway_pool) {
        if (strcmp(gw.host, host) == 0) {
            gw.latency = latency;
            // 通知网关状态更新
            if (gateway_callback) {
                gateway_callback(gw.host, gw.port, gw.latency, gateway_callback_userdata);
            }
            break;
        }
    }
}

void gateway_set_callback(gateway_change_cb_t cb, void* userdata) {
    gateway_callback = cb;
    gateway_callback_userdata = userdata;
}

Gateway* gateway_select() {
    if (gateway_pool.empty()) return nullptr;
    
    // 加权随机选择算法，考虑延迟因素
    uint32_t total_score = 0;
    for (auto& gw : gateway_pool) {
        uint32_t score = gw.weight;
        if (gw.latency > 0 && gw.latency < 100) {
            score *= 2; // 低延迟网关权重加倍
        } else if (gw.latency > 500) {
            score /= 2; // 高延迟网关权重减半
        }
        total_score += score;
    }
    
    uint32_t r = rand() % total_score;
    for (auto& gw : gateway_pool) {
        uint32_t score = gw.weight;
        if (gw.latency > 0 && gw.latency < 100) {
            score *= 2;
        } else if (gw.latency > 500) {
            score /= 2;
        }
        
        if (r < score) return &gw;
        r -= score;
    }
    
    return &gateway_pool[0];
}

void resource_register(const char* uri, const char* gateway_group) {
    resource_map[uri] = gateway_group;
}

const char* resource_locate(const char* uri) {
    auto it = resource_map.find(uri);
    return (it != resource_map.end()) ? it->second.c_str() : nullptr;
}

void input_packet(const uint8_t* data, size_t len) {
    if (g_packet_callback) {
        g_packet_callback(data, len, g_packet_callback_userdata);
    }
    if (stack_fd == 0) {
        return;
    }
    
    // 检测HTTP流量
    if (http_analysis_enabled && len > 4 && 
        memcmp(data, "GET ", 4) == 0 || memcmp(data, "POST ", 5) == 0) {
        HttpSessionStamp stamp = {0};
        if (parse_http_request(data, len, &stamp)) {
            stamp.timestamp = time(nullptr);
            stamp.signature = http_generate_signature((const char*)data, len);
            
            if (http_session_callback) {
                http_session_callback(&stamp, http_callback_userdata);
            }
        }
    }
    
    write(stack_fd, data, len);
    logger_info("zjx---input_packet-success");
}

static write_packet_t write_packet_fn = nullptr;

void socks_to_tun(int fd) {
    uint8_t buffer[BUFFER_SIZE];
    while (true) {
        ssize_t ret = read(fd, buffer, sizeof(buffer));
        if (ret <= 0) {
            if (errno == EAGAIN || errno == EWOULDBLOCK) {
                usleep(10000); // 10ms避免CPU忙等
                continue;
            }
            logger_info("zjx--socks read error: %{public}d", errno);
            break;
        }

        {
            std::lock_guard<std::mutex> lock(g_tun_mutex);
            if (write_packet_fn) {
                write_packet_fn(buffer, ret);
            } else {
                logger_info("zjx--write_packet_fn is null!");
            }
        }

        // 调试日志
        if (ret > 0 && buffer[0] != 0) {
            logger_info("zjx--Received %{public}zd bytes: [0]=%{public}02x", ret, buffer[0]);
        }
    }
    close(fd);
}
//void socks_to_tun(int fd) {
//    uint8_t buffer[BUFFER_SIZE] = {0};
//    while (1) {
//        int ret = read(fd, buffer, sizeof(buffer));
//        if (ret <= 0) {
//            if (errno != EAGAIN) {
//                sleep(1);
//            }
//            continue;
//        }
//        if (write_packet_fn != nullptr && ret > 0) {
//            write_packet_fn(buffer, ret);
//        }
//    }
//}

bool start_socks(write_packet_t write_packet, const char *proxy_host, int proxy_port) {
    int socket_pair[2];
    if (socketpair(AF_UNIX, SOCK_STREAM | SOCK_NONBLOCK, 0, socket_pair) == -1) {
        logger_info("zjx--socketpair failed: %{public}d", errno);
        return false;
    }

    {
        std::lock_guard<std::mutex> lock(g_tun_mutex);
        write_packet_fn = write_packet;
        stack_fd = socket_pair[1]; // 原子操作
    }

    std::thread([](int fd, const char* host, int port) {
        start_leaf(fd, host, port);
    }, socket_pair[0], proxy_host, proxy_port).detach();

    std::thread([](int fd) {
        socks_to_tun(fd);
    }, socket_pair[1]).detach();

    return true;
}
//bool start_socks(write_packet_t write_packet, const char *proxy_host, int proxy_port) {
//    int socket_pair[2];
//    if (socketpair(AF_UNIX, SOCK_STREAM, 0, socket_pair) == -1 ) { 
//        return false; 
//    }
//    write_packet_fn = write_packet;
//
//    std::thread(start_leaf, socket_pair[0], proxy_host, proxy_port).detach();
//    std::thread(socks_to_tun, socket_pair[1]).detach();
//    logger_info("zjx---start-tcp-success");
//    stack_fd = socket_pair[1];
//    return true;
//}
